import pandas as pd
import numpy as np
import os
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_absolute_error, mean_squared_error


def load_json_data(filename: str) -> pd.DataFrame:
    """加载JSON数据，处理路径问题和加载错误"""
    current_dir = os.path.dirname(os.path.abspath(__file__))
    file_path = os.path.join(current_dir, "..", "data", filename)
    file_path = os.path.abspath(file_path)

    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            data = pd.read_json(f)
        print(f"成功加载数据: {file_path}，共{len(data)}条记录")
        return data
    except FileNotFoundError:
        print(f"错误: 找不到数据文件 {file_path}")
        return None
    except Exception as e:
        print(f"加载数据时发生错误: {str(e)}")
        return None


# 加载课程需求数据
df = load_json_data("历年课程选课数据及影响因素.json")

# 检查数据加载是否成功
if df is None or df.empty:
    print("数据加载失败，无法继续执行程序")
    exit(1)

# 检查必要的列是否存在
required_columns = ['teacher_rating', 'course_difficulty', 'predicted_demand', 'actual_enrollment']
missing_columns = [col for col in required_columns if col not in df.columns]

if missing_columns:
    print(f"错误: 数据中缺少必要的列: {missing_columns}")
    print(f"数据中实际包含的列: {df.columns.tolist()}")
    exit(1)

# 特征工程 - 从学期提取年份和学期信息
df['year'] = df['semester'].apply(lambda x: int(x.split('-')[0]))
df['is_spring'] = df['semester'].apply(lambda x: 1 if 'Spring' in x else 0)

# 特征和目标变量
X = df[['teacher_rating', 'course_difficulty', 'year', 'is_spring']]
y = df['actual_enrollment']

# 数据集划分
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

# 训练模型
lr_model = LinearRegression()
rf_model = RandomForestRegressor(random_state=42, n_estimators=100)

lr_model.fit(X_train, y_train)
rf_model.fit(X_train, y_train)

# 预测
y_pred_lr = lr_model.predict(X_test)
y_pred_rf = rf_model.predict(X_test)

# 评估指标
metrics = {
    "linear_regression": {
        "MAE": round(mean_absolute_error(y_test, y_pred_lr), 3),
        "RMSE": round(np.sqrt(mean_squared_error(y_test, y_pred_lr)), 3),
        "R2": round(lr_model.score(X_test, y_test), 3)
    },
    "random_forest": {
        "MAE": round(mean_absolute_error(y_test, y_pred_rf), 3),
        "RMSE": round(np.sqrt(mean_squared_error(y_test, y_pred_rf)), 3),
        "R2": round(rf_model.score(X_test, y_test), 3)
    }
}

# 特征重要性
feature_importance = pd.DataFrame({
    "feature": X.columns,
    "rf_importance": rf_model.feature_importances_ if hasattr(rf_model, 'feature_importances_') else [0] * len(
        X.columns),
    "lr_coef": lr_model.coef_.tolist() if hasattr(lr_model, 'coef_') else [0] * len(X.columns)
}).sort_values("rf_importance", ascending=False)

# 生成未来课程数据
future_courses = pd.DataFrame({
    'teacher_rating': [4.5, 3.8, 4.2, 3.6, 4.7],
    'course_difficulty': [3.2, 4.1, 2.8, 3.5, 2.5],
    'year': [2024, 2024, 2025, 2025, 2025],
    'is_spring': [1, 0, 1, 0, 1]  # 1=Spring, 0=Fall
})

# 预测未来课程需求
future_pred_lr = np.round(lr_model.predict(future_courses)).astype(int)
future_pred_rf = np.round(rf_model.predict(future_courses)).astype(int)

# 转换为学期名称
future_courses['semester'] = future_courses.apply(
    lambda row: f"{row['year']}-Spring" if row['is_spring'] == 1 else f"{row['year']}-Fall", axis=1
)

# 封装结果
course_demand_result = {
    "historical_data": df[['semester', 'course_id', 'actual_enrollment']].to_dict("records"),
    "metrics": metrics,
    "feature_importance": feature_importance.to_dict("records"),
    "future_prediction": {
        "courses_info": future_courses[['semester', 'teacher_rating', 'course_difficulty']].to_dict("records"),
        "lr_prediction": future_pred_lr.tolist(),
        "rf_prediction": future_pred_rf.tolist()
    }
}

print("课程需求预测模型训练完成")
